/*
 * Code for blog.techhead.biz
 * Distributed under BSD-style license
 */

package biz.techhead.jdbc;

import biz.techhead.funcy.FuncE;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import javax.sql.DataSource;

/**
 *
 * @author jonathan
 */
public class DB {
    
    public static interface ResultSetMapperE<T,E extends Throwable>
            extends FuncE<T,ResultSet,Throwable> {
        T call(ResultSet rs) throws SQLException, E;
    }

    public static interface ResultSetMapper<T>
            extends ResultSetMapperE<T,RuntimeException> {
    }

    public static interface ConnectionActionE<T,E extends Throwable>
            extends FuncE<T,Connection,Throwable> {
        T call(Connection con) throws SQLException, E;
    }

    public static interface ConnectionAction<T>
            extends ConnectionActionE<T,RuntimeException> {
    }
    
    /**
     * Takes a ResultSet;
     * returns the first column of the first row as a String
     * or null if the ResultSet is empty
     */
    public static final ResultSetMapper<String>
            return1String = new ResultSetMapper<String>() {

        public String call(ResultSet rs) throws SQLException {
            return rs.next() ? rs.getString(1) : null;
        }
    };

    /**
     * Takes a ResultSet;
     * returns a Map<String,Object> representing the columns on the first row
     * or an empty Map if the ResultSet is empty.
     */
    public static final ResultSetMapper< Map<String,Object> >
            return1Map = new ResultSetMapper<Map<String,Object>>() {

        public Map<String,Object> call(ResultSet rs) throws SQLException {
            Map<String,Object> map = new HashMap<String,Object>();
            if (rs.next()) {
                ResultSetMetaData md = rs.getMetaData();
                for (int i=1,len=md.getColumnCount(); i<=len; i++) {
                    map.put(md.getColumnLabel(i), rs.getObject(i));
                }
            }
            return map;
        }
    };

    /**
     * Takes a ResultSet;
     * returns a List< Map<String,Object> > representing the entire ResultSet
     * as an in-memory Collection.
     */
    public static final ResultSetMapper< List<Map<String,Object>> >
            returnListOfMaps = new ResultSetMapper<List<Map<String,Object>>>() {

        public List<Map<String,Object>> call(ResultSet rs) throws SQLException {
            ArrayList<Map<String,Object>> list =
                    new ArrayList<Map<String,Object>>();
            ResultSetMetaData md = rs.getMetaData();
            String[] columnNames = new String[md.getColumnCount()];
            for (int i=0; i<columnNames.length; i++) {
                columnNames[i] = md.getColumnLabel(i+1);
            }
            while (rs.next()) {
                Map<String,Object> map = new HashMap<String,Object>();
                for (int i=0; i<columnNames.length; i++) {
                    map.put(columnNames[i], rs.getObject(i+1));
                }
                list.add(map);
            }
            return list;
        }
    };

    public DB(DataSource ds) {
        this.ds = ds;
    }

    public void setReuseConnection(boolean reuse) {
        reuseCon = reuse;
        if (!reuseCon) close();
    }

    @SuppressWarnings("unchecked")
    public <T,E extends Throwable>
            T execute(ConnectionActionE<T,E> ca) throws DBException, E {
        Throwable r = null;
        Connection con = null;
        try {
            con = getConnection();
            return ca.call(con);
        } catch (Throwable t) {
            r = t;
            if (t instanceof Error) throw (Error) t;
            if (t instanceof RuntimeException) throw (RuntimeException) t;
            if (t instanceof SQLException)
                throw new DBException(t.getMessage(), t);
            throw (E) t;
        } finally {
            close(con, r);
        }
    }

    public int executeUpdate(final String query) throws DBException {
        
        Integer r = execute( new ConnectionAction<Integer>() {
            Statement st;
            ResultSet rs;
            public Integer call(Connection con) throws SQLException {
                try {
                    st = con.createStatement();
                    st.execute(query, Statement.NO_GENERATED_KEYS);
                    return st.getUpdateCount();
                } finally {
                    close(rs);
                    close(st);
                }
            }
        } );
        
        return r;
    }

    public int executeUpdate(final String query, final Object... args)
            throws DBException {

        if (args.length == 0) return executeUpdate(query);

        Integer r = execute( new ConnectionAction<Integer>() {
            PreparedStatement st;
            ResultSet rs;
            public Integer call(Connection con) throws SQLException {
                try {
                    int i = 0;
                    st = con.prepareStatement(query,
                            Statement.NO_GENERATED_KEYS);
                    for (Object arg : args) {
                        st.setObject(++i, arg);
                    }
                    st.execute();
                    return st.getUpdateCount();
                } finally {
                    close(rs);
                    close(st);
                }
            }
        } );
        
        return r;
    }
    
    public <T,E extends Throwable>
            T executeUpdate(final ResultSetMapperE<T,E> p, final String query)
            throws DBException, E {
        
        if (p == null) {
            executeUpdate(query);
            return null;
        }

        T r = execute( new ConnectionActionE<T,E>() {
            Statement st;
            ResultSet rs;
            public T call(Connection con) throws SQLException, E {
                try {
                    st = con.createStatement();
                    st.execute(query, Statement.RETURN_GENERATED_KEYS);
                    rs = st.getGeneratedKeys();
                    return p.call(rs);
                } finally {
                    close(rs);
                    close(st);
                }
            }
        } );
        
        return r;
    }

    public <T,E extends Throwable>
            T executeUpdate(final ResultSetMapperE<T,E> p,
                            final String query,
                            final Object... args) throws DBException, E {

        if (p == null) {
            executeUpdate(query, args);
            return null;
        }

        if (args.length == 0) {
            return executeUpdate(p, query);
        }

        T r = execute( new ConnectionActionE<T,E>() {
            PreparedStatement st;
            ResultSet rs;
            public T call(Connection con) throws SQLException, E {
                try {
                    int i = 0;
                    st = con.prepareStatement(query,
                            Statement.RETURN_GENERATED_KEYS);
                    for (Object arg : args) {
                        st.setObject(++i, arg);
                    }
                    st.execute();
                    rs = st.getGeneratedKeys();
                    return p.call(rs);
                } finally {
                    close(rs);
                    close(st);
                }
            }
        } );
        
        return r;
    }

    public <T,E extends Throwable>
            T executeQuery(final ResultSetMapperE<T,E> p, final String query)
            throws DBException, E {

        T r = execute( new ConnectionActionE<T,E>() {
            Statement st;
            ResultSet rs;
            public T call(Connection con) throws SQLException, E {
                try {
                    st = con.createStatement();
                    rs = st.executeQuery(query);
                    return p.call(rs);
                } finally {
                    close(rs);
                    close(st);
                }
            }
        } );
        
        return r;
    }

    public <T,E extends Throwable>
            T executeQuery(final ResultSetMapperE<T,E> p,
                           final String query,
                           final Object... args) throws DBException, E {

        if (args.length == 0) { return executeQuery(p, query); }
        
        T r = execute( new ConnectionActionE<T,E>() {
            PreparedStatement st;
            ResultSet rs;
            public T call(Connection con) throws SQLException, E {
                try {
                    int i = 0;
                    st = con.prepareStatement(query);
                    for (Object arg : args) {
                        st.setObject(++i, arg);
                    }
                    rs = st.executeQuery();
                    return p.call(rs);
                } finally {
                    close(rs);
                    close(st);
                }
            }
        } );
        
        return r;
    }
    
    private Connection getConnection() throws SQLException {
        if (reuseCon) {
            if (lastCon == null)
                lastCon = ds.getConnection();
            return lastCon;
        }
        return ds.getConnection();
    }
    
    private void close(Statement st) {
        try { if (st != null) st.close(); }
        catch (Throwable t) { }
    }
    
    private void close(ResultSet rs) {
        try { if (rs != null) rs.close(); }
        catch (Throwable t) { }
    }

    private void close(Connection con, Throwable r) {
        if ( (con != null) && ( !reuseCon || (r != null) ) ) {
            close(con);
            if (con == lastCon) {
                lastCon = null;
            }
        }
    }

    private void close(Connection con) {
        try { con.close(); } catch (Throwable t) { }
    }
    
    public void close() {
        if (lastCon != null) {
            close(lastCon);
            lastCon = null;
        }
    }
    
    @Override
    protected void finalize() {
        close();
    }

    private DataSource ds;
    private boolean reuseCon;
    private Connection lastCon;
}
